import torch
import torch.nn as nn
from models.backbone.IREncoder import IREncoder
from models.head.decouplehead import DecoupleHead


class Aquarius(nn.Module):
    def __init__(self, num_det, num_seg, image_channels=3, radar_channels=3, resolution=320, image_encoder_width = [32, 48, 96, 176]):
        super(Aquarius, self).__init__()

        self.num_det = num_det
        self.num_seg = num_seg
        self.resolution = resolution

        self.image_channels = image_channels
        self.radar_channels = radar_channels

        self.image_radar_encoder = IREncoder(num_class_seg=num_seg, resolution=resolution, image_encoder_width=image_encoder_width)
        self.det_head = DecoupleHead(num_classes=num_det, image_encoder_width=image_encoder_width)

    def forward(self, x, x_radar):
        # print('in', x.shape, x_radar.shape) # [b, 3, h, h], [b, 3, h, h], h=320
        fpn_out, se_seg_output, lane_seg_output = self.image_radar_encoder(x, x_radar)
        det_output = self.det_head(fpn_out)
        # [b, 12, h/8, h/8], [b, 12, h/16, h/16], [b, 12, h/32, h/32], [b, 9, h, h], [b, 2, h, h]
        # print('out', det_output[0].shape, det_output[1].shape, det_output[2].shape, se_seg_output.shape, lane_seg_output.shape)
        return det_output, se_seg_output, lane_seg_output
